
import getopt,sys
import numpy as np
from scipy.integrate import odeint
import pickle as pickle


def haze_box_fn_moist(P_strat, P_moist,krain):

    vf = 0.03 # cm/s
    Ha = 10*1e5 # km to cm
    
    L = vf/Ha # cm/s x cm = s-1    
    Kzz = 1e5 # cm2/s 
    D = 1/(Ha*Ha/(2*Kzz)) # cm2 / cm2 /s  = s --> 1/s 

    dn_1 = lambda X: P_strat - X[0]*L + D*X[1] - X[0]*D
    dn_2 = lambda X: P_moist - X[1]*L - X[1]*krain + D*X[0] - 2*D*X[1]
        
    def dXdt(X,t):
        return [dn_1(X), dn_2(X)]

    nt = 1e3
    t_a = np.linspace(0,300*86400,int(nt))
    X0=[1,1]

    X_a = np.zeros((len(t_a),len(X0)))
    X_a[1,:] = X0
    X=X0
    
    solver_options = {'rtol': 1e-8}
    for it in range(1, int(nt)):
        X = odeint(dXdt, X, [t_a[it - 1], t_a[it]], hmax=1e4, rtol=solver_options['rtol'])[-1, :]
        X_a[it, :] = X

    X_f = X_a[-1,:]
    return X_f


def haze_box_fn_dry(P_strat):
    vf = 0.03 # cm/s
    Ha = 10*1e5 # km to cm
    L = vf/Ha # cm/s x cm = s-1
    
    #krain = 2.0e-6 # s-1
    Kzz = 1e5 # cm2/s
    D = 1/(Ha*Ha/(2*Kzz)) # cm2 / cm2 /s  = s --> 1/s
    
    dn_1 = lambda X: P_strat - X[0]*L - X[0]*D

    def dXdt(X,t):
        return [dn_1(X)]

    nt = 1e3
    t_a = np.linspace(0,300*86400,int(nt))
    X0=[1]

    X_a = np.zeros((len(t_a),len(X0)))
    X_a[1,:] = X0
    X=X0

    solver_options = {'rtol': 1e-8}
    for it in range(1, int(nt)):
        X = odeint(dXdt, X, [t_a[it - 1], t_a[it]], hmax=1e4, rtol=solver_options['rtol'])[-1, :]
        X_a[it, :] = X

    X_f = X_a[-1,:]
    return X_f


infile = open('../sulfur_box/tau_rain.txt','r')
temp = infile.readline().split()
kfv,kTs,k = [],[],[]
while len(temp)>0:
    kfv += [float(temp[0])]
    kTs += [float(temp[1])]
    k += [float(temp[2])]
    temp = infile.readline().split()

fname = 'summary_kineticsresults_eddy1e5_shield300_krON_jan25.txt'
outfile = open('hazetot_eddy1e5_shield300_krON_jan25.txt','w')
infile = open(fname,'r')
temp = infile.readline().split()
hazestrat = []
hazemoist = []
Tsurfs = []
waters = []
prodstratL = []
prodrainL = []
while len(temp)>1:
    fv = float(temp[1])
    Ts = float(temp[0])
    pstrat = float(temp[2])
    prain = float(temp[3])

    try:
        indx = np.where((np.array(kfv) == fv) & (np.array(kTs)==Ts))[0][0]
        new_value = k[indx]
    except:
        new_value = 0.0
        print(fv,Ts)
    
    
    Tsurfs += [Ts]
    waters += [fv]
    prodstratL+=[pstrat]
    prodrainL += [prain]
    
    rho = 1.8 
    r = 0.50*1e-4
    wtpct = 0.75
    mu= 98.0 
    mp = 1.67e-24 
    mpart = rho*(4/3.0)*3.141592*r*r*r
    prodstrat = pstrat*mp*mu/(mpart*wtpct)
    prodrain = prain*mp*mu/(mpart*wtpct)


    if prodrain >0:
        Xa = haze_box_fn_moist(prodstrat,prodrain,new_value)
        temp = '%.3e'%Ts+'  '+'%.3e'%fv+'  '
        for i in Xa:
            temp+='%.3e'%i+'  '
        outfile.write(temp+'\n')
        hazestrat+=[float(Xa[0])]
        hazemoist += [float(Xa[1])]
    else:
        Xa = haze_box_fn_dry(prodstrat)
        temp = '%.3e'%Ts+'  '+'%.3e'%fv+'  '+'%.3e'%Xa[0]+'  '+'%.3e'%0+'\n'
        hazestrat+=[float(Xa[0])]
        hazemoist+=[0.0]
        outfile.write(temp)

    temp = infile.readline().split()

dat = {'prod_strat':prodstratL,'prod_moist':prodrainL,'Tsurf':Tsurfs,'fv':waters,'haze_strat':hazestrat,'haze_moist':hazemoist}
with open('haze_krain_1e5.pickle','wb') as handle:
    pickle.dump(dat,handle,protocol=pickle.HIGHEST_PROTOCOL)
    
